import numpy as np
import matplotlib.pyplot as plt


def swl_1(w, ceita=0.04, ld=6, cd=1, fz=1):
    """
    # 地震波设计谱函数
    :param w: 圆频率
    :param ceita: 阻尼比
    :param ld: 地震烈度
    :param cd: 场地类型
    :param fz: 场地分组
    :return:
    """
    # 烈度选择，影响最大值
    if ld == 6:
        arfmax = 0.11
    if ld == 7:
        arfmax = 0.23
    if ld == 8:
        arfmax = 0.45
    if ld == 9:
        arfmax = 0.90
    #  场地类别，设计地震分组选择，影响特征周期Tg
    if cd == 1:
        if fz == 1:
            Tg = 0.25
        if fz == 2:
            Tg = 0.30
        if fz == 3:
            Tg = 0.35
    if cd == 2:
        if fz == 1:
            Tg = 0.35
        if fz == 2:
            Tg = 0.40
        if fz == 3:
            Tg = 0.45
    if cd == 3:
        if fz == 1:
            Tg = 0.45
        if fz == 2:
            Tg = 0.55
        if fz == 3:
            Tg = 0.65
    if cd == 4:
        if fz == 1:
            Tg = 0.65
        if fz == 2:
            Tg = 0.75
        if fz == 3:
            Tg = 0.90

    def designrs(ww):
        lmt1 = 0.02 + (0.05 - ceita) / 8
        if lmt1 < 0:
            lmt1 = 0
        lmt2 = 1 + (0.05 - ceita) / (0.06 + 1.7 * ceita)
        if lmt2 < 0.55:
            lmt2 = 0.55
        sjzs = 0.9 + (0.05 - ceita) / (0.5 + 5 * ceita)
        # 分段位置 T1 T2 T3
        T1 = 0.1
        T2 = Tg  # 特征周期
        T3 = 5 * Tg
        T_jg = 2 * np.pi / ww  # 结构特征周期
        # 第一段 0～T1
        if T_jg <= T1:
            arf_jg = 0.45 * arfmax + (lmt2 * arfmax - 0.45 * arfmax) / 0.1 * T_jg
        # 第二段 T1～T2
        if (T1 < T_jg) & (T_jg <= T2):
            arf_jg = lmt2 * arfmax
        # 第三段 T2~T3
        if (T2 < T_jg) & (T_jg <= T3):
            arf_jg = ((Tg / T_jg) ** sjzs) * lmt2 * arfmax
        # 第四段 T3～6.0
        if (T3 < T_jg) & (T_jg <= 6.0):
            arf_jg = (lmt2 * 0.2 ** sjzs - lmt1 * (T_jg - 5 * Tg)) * arfmax
        # 第五段 6.0～
        if 6.0 < T_jg:
            arf_jg = (lmt2 * 0.2 ** sjzs - lmt1 * (6.0 - 5 * Tg)) * arfmax
        # 反应谱值 拟加速度值
        rs_z = arf_jg * 9.8
        return rs_z

    p = 0.9  # 反应超越概率
    st = designrs(w)
    sw = ceita / (np.pi * w) * st ** 2 / np.log(
        -np.pi / (w * (2 * np.pi / w) * np.sqrt(1 - ceita ** 2)) * np.log(1 - p))
    return sw


def swl_2(category, ld, w):  # 输入场地类型，地面运动加速度峰值，频率
    # Kanai-Tajimi 功率谱
    wgF = 15.6
    wgM = 10.4
    wgS = 5.2
    rF = 3
    rM = 2.9
    rS = 2.7
    g = 9.8
    if ld == 6:
        arfmax = 0.11
    if ld == 7:
        arfmax = 0.23
    if ld == 8:
        arfmax = 0.45
    if ld == 9:
        arfmax = 0.90
    PGA = arfmax * g
    # 场地类型  1为硬  2为中等  3为软
    if category == 1:
        wg = wgF
        r = rF
    if category == 2:
        wg = wgM
        r = rM
    if category == 3:
        wg = wgS
        r = rS

    #  计算功率谱密度
    ksig = wg / 25
    S0 = 2 * ksig * PGA ** 2 / (np.pi * r ** 2 * wg * (1 + 4 * ksig ** 2))
    sw = ((1 + 4 * ksig ** 2 * (w / wg) ** 2) / ((1 - (w / wg) ** 2) ** 2 + 4 * ksig ** 2 * (w / wg) ** 2)) * S0
    return sw


def swl_3(category, ld, w):  # 输入场地类型，地面运动加速度峰值，频率
    # Ruiz-Penzien修正的Kanai-Tajimi 功率谱
    wgF = 15.6
    wgM = 10.4
    wgS = 5.2
    rF = 3
    rM = 2.9
    rS = 2.7
    g = 9.8
    if ld == 6:
        arfmax = 0.11
    if ld == 7:
        arfmax = 0.23
    if ld == 8:
        arfmax = 0.45
    if ld == 9:
        arfmax = 0.90
    PGA = 2.25 * arfmax * g
    # 场地类型  1为硬  2为中等  3为软
    if category == 1:
        wg = wgF
        r = rF
    if category == 2:
        wg = wgM
        r = rM
    if category == 3:
        wg = wgS
        r = rS
    #  计算功率谱密度
    ksig = wg / 25
    wf = wg / 10
    ksif = 0.6
    S0 = 2 * ksig * PGA ** 2 / (np.pi * r ** 2 * wg * (1 + 4 * ksig ** 2))
    sw = ((1 + 4 * ksig ** 2 * (w / wg) ** 2) / ((1 - (w / wg) ** 2) ** 2 + 4 * ksig ** 2 * (w / wg) ** 2)) * (
            (w / wf) ** 4 / ((1 - (w / wf) ** 2) ** 2 + 4 * (w / wf) ** 2 * ksif ** 2)) * S0
    return sw


def swl_4(category, ld, w):
    # 由欧进萍等人修正的改进金井清模型
    wgF = 15.6
    wgM = 10.4
    wgS = 5.2
    rF = 3
    rM = 2.9
    rS = 2.7
    g = 9.8
    if ld == 6:
        arfmax = 0.11
    if ld == 7:
        arfmax = 0.23
    if ld == 8:
        arfmax = 0.45
    if ld == 9:
        arfmax = 0.90
    PGA = 2.25 * arfmax * g
    # 场地类型  1为硬  2为中等  3为软
    if category == 1:
        wg = wgF
        r = rF
    if category == 2:
        wg = wgM
        r = rM
    if category == 3:
        wg = wgS
        r = rS
    #  计算功率谱密度
    ksig = wg / 25
    S0 = 2 * ksig * PGA ** 2 / (np.pi * r ** 2 * wg * (1 + 4 * ksig ** 2))
    sw = ((1 + 4 * ksig ** 2 * (w / wg) ** 2) / ((1 - (w / wg) ** 2) ** 2 + 4 * ksig ** 2 * (w / wg) ** 2)) \
         * (1 / (1 + w ** 2 / (8 * np.pi) ** 2)) * S0
    return sw


def swl_5(category, ld, w):
    # 李春祥修正的Ruiz-Penzien功率谱
    wgF = 15.6
    wgM = 10.4
    wgS = 5.2
    rF = 3
    rM = 2.9
    rS = 2.7
    g = 9.8
    if ld == 6:
        arfmax = 0.11
    if ld == 7:
        arfmax = 0.23
    if ld == 8:
        arfmax = 0.45
    if ld == 9:
        arfmax = 0.90
    PGA = 2.25 * arfmax * g
    # 场地类型  1为硬  2为中等  3为软
    if category == 1:
        wg = wgF
        r = rF
    if category == 2:
        wg = wgM
        r = rM
    if category == 3:
        wg = wgS
        r = rS
    #  计算功率谱密度
    ksig = wg / 25
    wf = wg / 10
    ksif = 0.6
    S0 = 2 * ksig * PGA ** 2 / (np.pi * r ** 2 * wg * (1 + 4 * ksig ** 2))
    sw = ((1 + 4 * ksig ** 2 * (w / wg) ** 2) / ((1 - (w / wg) ** 2) ** 2 + 4 * ksig ** 2 * (w / wg) ** 2)) \
         * ((w / wf) ** 4 / ((1 - (w / wf) ** 2) ** 2 + 4 * (w / wf) ** 2 * ksif ** 2)) * (
                 1 / (1 + w ** 2 / (8 * np.pi) ** 2)) * S0
    return sw


# 测试
if __name__ == '__main__':
    a = np.arange(0.01, 80, 0.01)  # 频率序列
    b = np.zeros_like(a)  # 功率谱密度序列
    for i in range(len(a)):
        b[i] = swl_1(a[i], 0.04, 8, 3, 2)
    # for i in range(len(a)):
    #     b[i] = swl_5(1, 6, a[i])
    print(b)
    plt.plot(a, b)
    plt.show()
